"""
    ================================================================================
                            ------------utf-8--------------
    ================================================================================
@Author: 
    rfdsg
@Create Time: 
    2024/10/18 - 17:14
@Description:

@Attention:
    
"""

from typing import Callable, Generator, Literal
import numpy as np
import pandas as pd
from typing_extensions import Union, Optional

from factor_lib.common_tool import DSN, DS
from factor_lib.common_tool.easy_type import DSNL


class TransInfo:
    def __init__(self,
                 *,
                 factor: Optional[Union[Generator, DSNL]],
                 columns: Optional[pd.Index],
                 returns_for_position: Optional[DSN] = None,
                 y: Optional[DS] = None,
                 price: Optional[DS] = None,
                 deal_index: Optional[Union[Generator, DSNL]] = None,
                 market: Optional[pd.Series] = None,
                 select_mask: Optional[np.ndarray] = None,
                 time: pd.Series = None,
                 commission: int = 0.0005
                 ):
        """

        Args:
            factor: 因子
            returns_for_position: 计算仓位净值的returns， 包含代码名称, 必须是明天比今天存今天 - 1 的returns
            y: 目标列,是以1.~为类型的
            price: 标的价格
            market: 市场基准标的收益率
        """
        self.columns = columns
        self.__factor = factor
        self.__y = y
        self.__total_time = time
        self.time = self.__total_time
        self.returns_for_position = returns_for_position
        self.market = market
        self.net_returns = price
        self.select_mask = None
        self.best_decay = []
        if select_mask is not None:
            self.select_mask = select_mask
            self.__factor = self.__factor[select_mask]
            self.__y = self.__y[select_mask]
            self.time = self.__total_time[select_mask]

        self.__deal_index = deal_index
        self.__position = np.array([])
        self.__net = np.array([])
        self.model = None
        self.back = {}
        self.fun_str = None
        self.other_info = {}
        self.commission = commission
        self.factor_len = None
        self.deal_index_len = None


    @property
    def factor(self):
        return self.__factor.copy()

    @property
    def total_time(self):
        return self.__total_time.copy()

    @property
    def y(self):
        if self.__y is None:
            return None
        return self.__y.copy()

    @property
    def net(self):
        return self.__net.copy()

    @net.setter
    def net(self, net):
        if isinstance(net, (np.ndarray, list)):
            self.__net = net
        else:
            raise ValueError('错误类型')

    @property
    def deal_index(self):
        return self.__deal_index.copy()

    @deal_index.setter
    def deal_index(self, deal_index):
        if isinstance(deal_index, (pd.DataFrame, pd.Series, np.ndarray)):
            self.__deal_index = deal_index
        else:
            raise ValueError('错误类型')

    @property
    def position(self):
        return self.__position.copy()

    @position.setter
    def position(self, position):
        if isinstance(position, (pd.DataFrame, pd.Series, np.ndarray, list)):
            self.__position = position
        else:
            raise ValueError('错误类型')

    def map_back(self, data: Union[np.ndarray, pd.Series]):
        if self.select_mask is not None:
            target = self.total_time
            target[self.select_mask] = data
            target[~self.select_mask] = np.nan
            target.ffill(inplace=True)
        else:
            target = data
        return target

    def align_drop(self):
        """

        Returns:
            cal_return: 用以后续的position的net计算


        """
        merge_data = pd.concat([self.__factor, self.returns_for_position], axis=1)
        merge_data.dropna(axis=0, inplace=True)
        self.__factor = merge_data.iloc[:, :-1]
        self.returns_for_position = merge_data.iloc[:, -1]

    def get_cal(self,
                back: pd.DataFrame,
                deal_index: Union[pd.Series, np.ndarray],
                dynamic_position: np.ndarray,
                net: np.ndarray,
                model: Optional[Callable] = None
                ):
        self.back = back
        self.deal_index = deal_index
        self.__position = dynamic_position
        self.net = net
        self.model = model

    def factor_replace(self):
        # 直接替换
        for i, replacement in enumerate(self.__factor.columns):
            self.fun_str = self.fun_str.replace(f'X{i}', replacement)

    def recall_len(self, name: Literal['factor', 'deal_index'], lens: int):
        if name == 'factor':
            self.factor_len = lens
        elif name == 'deal_index':
            self.deal_index_len = lens
        else:
            raise TypeError('输入类型错误')

    def __getstate__(self):
        # 保留的属性集合
        keep_attrs = {'program', 'back', 'factor_name', 'code', 'fun_str'}

        # 检查所有必需的属性是否存在，确保健壮性
        missing_attrs = keep_attrs - self.__dict__.keys()
        if missing_attrs:
            raise AttributeError(f"Missing required attributes: {missing_attrs}")

        # 返回存在的属性
        return {k: v for k, v in self.__dict__.items() if k in keep_attrs}

    def __setstate__(self, state):
        # 恢复状态
        self.__dict__.update(state)

    def __reduce__(self):
        # 自定义如何保存对象及其状态
        state = self.__getstate__()
        return self._rebuild_instance, (state,)

    @staticmethod
    def _rebuild_instance(state):
        # 自定义反序列化逻辑，重建对象
        obj = TransInfo.__new__(TransInfo)  # 创建空对象，不调用 __init__
        obj.__setstate__(state)  # 还原状态
        return obj


def several_to_single(transinfo_list: [TransInfo]) -> TransInfo:
    for transinfo in transinfo_list:
        if isinstance(transinfo, TransInfo):
            return transinfo


class DataEngine:
    def __init__(self):
        """
        初始化数据引擎时不加载任何数据，所有数据通过懒加载方式存储。
        """
        self._data_cache = dict({})  # 用字典来缓存数据
        self.__allowed_key = {
            'factor'
            'returns_for_position',
            'y',
            'net_returns',
            'deal_index',
            'market',
            'select_mask',
            'back',
            'dynamic_position',
            'net',
            'model',
            'raw_length'
        }

    def __call__(self, data_dict: dict):
        """
        更新数据。可以从外部输入新的数据并更新缓存。
        """
        keys = set(data_dict.keys())
        without_allowed_keys = keys - self.__allowed_key
        allowed_keys = keys & self.__allowed_key
        allowed_data = {k: data_dict[k] for k in data_dict.keys() & allowed_keys}
        self._data_cache.update(allowed_data)
        print(f"Data for {allowed_data.keys()} updated.")
        return without_allowed_keys

    def get_all_data_keys(self):
        """
        获取当前所有数据的键
        """
        return list(self._data_cache.keys())
